#!/usr/bin/env python
# -*- coding: utf-8 -*-

""" By Martin Senande-Rivera
	For Spatial and temporal expansion of global wildland fire activity in response to climate change """

import os
import glob, sys
import numpy as np
import pandas as pd
from pandas import Series, DataFrame
import xarray as xr
import datetime as dt 

ruta_in='./WFDE5/DATA/'
ruta_out='./WFDE5/'

## Precipitation
tps=sorted(glob.glob(ruta_in+'Rainf_WFDE5_CRU+GPCC_*.nc'))
start_date = dt.datetime.strptime('1996-01', '%Y-%m')
dates = pd.date_range(start_date, periods=len(tps), freq='m')

lons = np.linspace(-179.75,179.75,720)
lats = np.linspace(-89.75,89.75,360)
TP = xr.DataArray(np.zeros([len(tps),360, 720]), 
	dims=['time','lat', 'lon'], 
	coords={'time':dates, 'lat': lats, 'lon': lons},
	name='tp',
	attrs={'units':'mm'})

for i in range(len(tps)):
	ds_tp = xr.open_dataset(tps[i])
	TP.values[i,:,:] = ds_tp.sum(dim='time').Rainf*3600.

TP.to_netcdf(ruta_out+'wfde5_tp_months.nc')

ds_tp_m = TP.groupby('time.month').mean(dim='time')
ds_tp_a = TP.groupby('time.year').sum(dim='time').mean(dim='year')

ds_tp_m.to_netcdf(ruta_out+'wfde5_tp_m_mean.nc')
ds_tp_a.to_netcdf(ruta_out+'wfde5_tp_a_mean.nc')

## Snow
tps=sorted(glob.glob(ruta_in+'Snowf_WFDE5_CRU+GPCC_*.nc'))
start_date = dt.datetime.strptime('1996-01', '%Y-%m')
dates = pd.date_range(start_date, periods=len(tps), freq='m')

lons = np.linspace(-179.75,179.75,720)
lats = np.linspace(-89.75,89.75,360)
TP = xr.DataArray(np.zeros([len(tps),360, 720]), 
	dims=['time','lat', 'lon'], 
	coords={'time':dates, 'lat': lats, 'lon': lons},
	name='tp',
	attrs={'units':'mm'})

for i in range(len(tps)):
	ds_tp = xr.open_dataset(tps[i])
	TP.values[i,:,:] = ds_tp.sum(dim='time').Snowf*3600.

TP.to_netcdf(ruta_out+'wfde5_sn_months.nc')

ds_tp_m = TP.groupby('time.month').mean(dim='time')
ds_tp_a = TP.groupby('time.year').sum(dim='time').mean(dim='year')

ds_tp_m.to_netcdf(ruta_out+'wfde5_sn_m_mean.nc')
ds_tp_a.to_netcdf(ruta_out+'wfde5_sn_a_mean.nc')


## Temperature
t2ms=sorted(glob.glob(ruta_in+'Tair_WFDE5_CRU_*.nc'))
start_date = dt.datetime.strptime('1996-01', '%Y-%m')
dates = pd.date_range(start_date, periods=len(t2ms), freq='m')

lons = np.linspace(-179.75,179.75,720)
lats = np.linspace(-89.75,89.75,360)
T2M = xr.DataArray(np.zeros([len(t2ms),360, 720]), 
	dims=['time','lat', 'lon'], 
	coords={'time':dates, 'lat': lats, 'lon': lons},
	name='t2m',
	attrs={'units':'Celsius'})

for i in range(len(t2ms)):
	ds_t2m = xr.open_dataset(t2ms[i])
	T2M.values[i,:,:] = ds_t2m.mean(dim='time').Tair-273.15

T2M.to_netcdf(ruta_out+'wfde5_t2m_months.nc')

ds_t2m_m = T2M.groupby('time.month').mean(dim='time')
ds_t2m_a = T2M.groupby('time.year').mean(dim='time').mean(dim='year')

ds_t2m_m.to_netcdf(ruta_out+'wfde5_t2m_m_mean.nc')
ds_t2m_a.to_netcdf(ruta_out+'wfde5_t2m_a_mean.nc')

